package dan.frontend
import chisel3._
import chisel3.util._
import dan.common.CoreModule
import dan.common.CoreBundle
import chisel3.SpecifiedDirection.Flip
import dan.utility.WrapInc
import dan.utility.WrapAdd
import dan.utility.MaskLower
import dan.common.Consts.CFI_BL
import dan.common.Consts.CFI_JIRL

class FTQIO extends CoreBundle with HasFTQParam{
    val enq = Flipped(Decoupled(new FetchPack()))
    val enqIdx = Output(UInt(idxBits.W))
    val deq =  Flipped(Valid(UInt(idxBits.W)))
    val getEntryPC = Vec(2, new FTQIdx2PC())
    val redirectIdx = Input(Valid(UInt(idxBits.W)))
    // 后端发来的分支预测更新
    val brUpdateInfo = Input(new BrUpdateInfo())
    // 传给分支预测器的更新信息
    val bpuUpdate = Output(Valid(new PredictionUpdate()))
    val rasUpdateEn = Output(Bool())
    val rasUpdateIdx = Output(UInt(log2Ceil(rasSize).W))
    val rasUpdateTarget = Output(UInt(targetBits.W))
}

class FTQ extends CoreModule with HasFTQParam{
    val io = IO(new FTQIO())
    val trainPtr = RegInit(0.U(idxBits.W))
    val commitPtr = RegInit(0.U(idxBits.W))
    val nextPtr = RegInit(1.U(idxBits.W))
    val full = (WrapInc(nextPtr, ftqSize) === trainPtr) || WrapAdd(nextPtr, 2.U, ftqSize) === trainPtr
    
    val pcs = Reg(Vec(ftqSize, UInt(vaBits.W)))
    val meta = SyncReadMem(ftqSize, Vec(fetchWidth, new BPUMeta()))
    val ftq = Reg(Vec(ftqSize, new FTQPack()))
    val rasPtrs = SyncReadMem(ftqSize, new RASIdx())
    val preRASPtr = RegInit(0.U.asTypeOf(new RASIdx()))
    val preFTQEntry = RegInit(0.U.asTypeOf(new FTQPack()))
    val prePC = RegInit(0.U(vaBits.W))

    val newEntry = Wire(new FTQPack())
    newEntry.cfiIdx := io.enq.bits.cfiIdx
    newEntry.cfiTaken := io.enq.bits.cfiIdx.valid
    newEntry.cfiMispredict := false.B
    newEntry.cfiType := io.enq.bits.cfiType
    newEntry.cfiIsCall := io.enq.bits.cfiIsCall
    newEntry.cfiIsRet := io.enq.bits.cfiIsRet
    newEntry.rasTopEntry := io.enq.bits.rasTopEntry
    newEntry.rasIdx := io.enq.bits.rasPtr.bits
    newEntry.brMask := io.enq.bits.brMask & io.enq.bits.mask
    val newRASPtr = preRASPtr.update(preFTQEntry.cfiIdx.valid, preFTQEntry.cfiIsCall, preFTQEntry.cfiIsRet)
    // 入队
    when(io.enq.fire){
        pcs(nextPtr) := io.enq.bits.pc
        meta.write(nextPtr, io.enq.bits.bpuMeta)
        ftq(nextPtr) := newEntry
        rasPtrs.write(nextPtr, newRASPtr)
        nextPtr := WrapInc(nextPtr, ftqSize)
        // 更新pre
        prePC := io.enq.bits.pc
        preFTQEntry := newEntry
        preRASPtr := newRASPtr
    }
    io.enqIdx := nextPtr
    when(io.deq.valid){
        commitPtr := io.deq.bits
    }
    // 正常情况下等于trainPtr，发生重定向时等于导致redirect的ftq索引
    // 确定用于训练bpu的信息
    val trainIdx = Mux(io.redirectIdx.valid, io.redirectIdx.bits, trainPtr)
    val trainEntry = RegNext(ftq(trainIdx))
    val trainRASPtr = rasPtrs.read(trainIdx)
    val trainMeta = meta.read(trainIdx, true.B)
    val trainPC = RegNext(pcs(trainIdx))
    val trainTarget = RegNext(pcs(WrapInc(trainIdx, ftqSize)))
    // train -> commit -> next
    //    | _ _ _ _ ... _ _ |
    val trainEnable = (trainPtr =/= commitPtr) && (WrapInc(trainPtr, ftqSize) =/= nextPtr) && 
        (!io.redirectIdx.valid && !RegNext(io.redirectIdx.valid))
    io.bpuUpdate.valid := false.B
    io.bpuUpdate.bits := DontCare
    when(RegNext(trainEnable)){
        val cfiIdx = trainEntry.cfiIdx.bits
        io.bpuUpdate.valid := trainEntry.cfiIdx.valid || trainEntry.brMask =/= 0.U
        io.bpuUpdate.bits.pc := trainPC
        // 屏蔽cfiIdx之后的所有指令
        val x = MaskLower(UIntToOH(cfiIdx)) & trainEntry.brMask
        io.bpuUpdate.bits.brMask := Mux(trainEntry.cfiIdx.valid, x, trainEntry.brMask)
        io.bpuUpdate.bits.idx := trainEntry.cfiIdx
        io.bpuUpdate.bits.mispredict := trainEntry.cfiMispredict
        io.bpuUpdate.bits.taken := trainEntry.cfiTaken
        io.bpuUpdate.bits.target := trainTarget
        io.bpuUpdate.bits.isBr := trainEntry.brMask(cfiIdx)
        io.bpuUpdate.bits.isJal := trainEntry.cfiType === CFI_BL || trainEntry.cfiType === CFI_JIRL
        io.bpuUpdate.bits.meta := trainMeta
    }
    when(trainEnable){
        trainPtr := WrapInc(trainPtr, ftqSize)
    }
    io.enq.ready := RegNext(!full || trainEnable)
    // 重定向
    // 入队指针指向导致重定向的指令的下一个位置，覆盖错误路径
    when(io.redirectIdx.valid){
        nextPtr := WrapInc(io.redirectIdx.bits, ftqSize)
    }
    val rasUpdateEn = WireInit(false.B)
    val rasUpdateTarget = WireInit(0.U(targetBits.W))
    val rasUpdateIdx = WireInit(0.U(log2Ceil(rasSize).W))
    
    val badEntry = ftq(io.redirectIdx.bits)
    val alterEntry = WireInit(badEntry)
    when(io.redirectIdx.valid){
        when(io.brUpdateInfo.resolInfo.mispredict){
            val newCfiIdx = io.brUpdateInfo.resolInfo.uop.pcLow(log2Ceil(fetchWidth)+log2Ceil(instrBytes)-1, log2Ceil(instrBytes))
            alterEntry.cfiIdx.valid := true.B
            alterEntry.cfiIdx.bits := newCfiIdx
            alterEntry.cfiTaken := io.brUpdateInfo.resolInfo.taken
            alterEntry.cfiMispredict := true.B
            alterEntry.cfiIsCall := badEntry.cfiIsCall && badEntry.cfiIdx.bits === newCfiIdx
            alterEntry.cfiIsRet := badEntry.cfiIsRet  && badEntry.cfiIdx.bits === newCfiIdx
        }
        rasUpdateEn := true.B
        rasUpdateTarget := badEntry.rasTopEntry
        rasUpdateIdx := badEntry.rasIdx
    }.elsewhen(RegNext(io.redirectIdx.valid)){
        preFTQEntry := RegNext(alterEntry)
        preRASPtr := trainRASPtr
        prePC := trainPC
        ftq(RegNext(io.redirectIdx.bits)) := RegNext(alterEntry)
    }
    
    io.rasUpdateEn := RegNext(rasUpdateEn)
    io.rasUpdateTarget := RegNext(rasUpdateTarget)
    io.rasUpdateIdx := RegNext(rasUpdateIdx)
    // FTQIdx2PC 读
    for(i <- 0 until 2){
        val idx = io.getEntryPC(i).idx
        val nextIdx = WrapInc(idx, ftqSize)
        val nextIsEnq = (nextIdx === nextPtr) && io.enq.fire
        val nextPC = Mux(nextIsEnq, io.enq.bits.pc, pcs(nextIdx))
        io.getEntryPC(i).entry := RegNext(ftq(idx))
        if (i == 0) {
            io.getEntryPC(i).pc := RegNext(pcs(idx))
            io.getEntryPC(i).rasPtr := DontCare
            io.getEntryPC(i).commitPC := RegNext(pcs(Mux(io.deq.valid, io.deq.bits, commitPtr)))
            io.getEntryPC(i).nextPCValid := RegNext(nextIdx =/= nextPtr || nextIsEnq)
            io.getEntryPC(i).nextPC := RegNext(nextPC)
        } else {
            //pcs fetch at s1 for better timing 
            io.getEntryPC(i).pc := pcs(RegNext(idx))
            io.getEntryPC(i).rasPtr := rasPtrs.read(idx)
            io.getEntryPC(i).commitPC := DontCare
            io.getEntryPC(i).nextPCValid := DontCare
            io.getEntryPC(i).nextPC := DontCare
        }
    }
} 

